
#!/usr/bin/env python


import rospy
from trajectory_msgs.msg import MultiDOFJointTrajectory

import numpy as np
import acado
import time
import matplotlib.pyplot as plt
from numpy import sin, cos,exp
from tf.transformations import euler_from_quaternion, quaternion_from_euler
from ros_numpy import numpify,msgify
from visualization_msgs.msg import MarkerArray,Marker
import time
from trajectory_msgs.msg import MultiDOFJointTrajectory
from sensor_msgs.msg import PointCloud2
from geometry_msgs.msg import Pose,PoseStamped,PoseArray

import open3d as o3d

pc_size=30

N=20 #timestep
NU=2
NOD=pc_size*3+3
NX=4
NY=6
NYN=4

NW=6
NWN=4


def pc2_msg_to_np(pc_msg):
    pc_np=numpify(pc_msg)
    pc_np= np.stack([pc_np[f] for f in ('x', 'y', 'z')])

    return pc_np

def np_to_pc2_msg(pc_np):
    data=np.zeros(pc_np.shape[1],
    dtype=[('x',np.float32), ('y',np.float32), ('z',np.float32)])
    data['x']=pc_np[0,]
    data['y']=pc_np[1,]
    data['z']=pc_np[2,]

    pc_msg=msgify(PointCloud2,data)
    pc_msg.is_dense=False
        
    return pc_msg

class ViewPlanning:
    def __init__(self):
        
        
        rospy.init_node('view_planning_node')
        
        self.pc_W=np.ones([3,pc_size])
        self.target_position=np.zeros([3,1])
        
        self.pc_pub=rospy.Publisher('/slot/downsampled_pointcloud',PointCloud2,queue_size=10)
        self.pose_pub=rospy.Publisher('/slot/pose_array',PoseArray,queue_size=10)

        #rospy.Subscriber('/position_traj', MultiDOFJointTrajectory, multi_dof_trajectory_callback)  # Subscribe to the "MultiDOFJointTrajectory" topic
        rospy.Subscriber('/slot/pointcloud',PointCloud2,self.pc_callback)
        rospy.Subscriber('/slot/target/pose', PoseStamped, self.target_callback)
        rospy.Subscriber('/position_traj', MultiDOFJointTrajectory, self.multi_dof_trajectory_callback) 
    
    def target_callback(self,target_pose_msg):
        pose_np=numpify(target_pose_msg.pose)
        self.target_position=pose_np[0:3,3]
    
    def multi_dof_trajectory_callback(self,msg):
        # Process the multi DOF joint trajectory message
        # Extract and print the poses and twists from the message

        new_msg=msg
        pc_W=self.pc_W
        
        x0=np.zeros(NX).reshape(1,NX)
        od=np.zeros(NOD)
        y=np.zeros([N,NY])
        u=np.zeros([N,NU])
        Q=np.eye(NW)
        Q[2,2]=10
        Q[3,3]=10
        Q=np.diag([0.,0.,10.,10.,0.1,0.1])
        W=np.transpose(np.tile(Q,N))
        WN=Q[:-2,:-2]
        X=np.tile(x0.reshape(NX,1),N+1).T
        OD=np.ones([N+1,NOD])

        OD[:,pc_size*0:pc_size*1]=pc_W[0,:]  #x
        OD[:,pc_size*1:pc_size*2]=pc_W[1,:]  #y
        OD[:,pc_size*2:pc_size*3]=pc_W[2,:]  #z
        
        xlog=[]
        
        marker_array=MarkerArray()
        pose_array=PoseArray()
        pose_array.header.frame_id="world"
        for p in range(len(msg.points)):
            X,u=acado.mpc(0,1,x0,X,u,OD,y,W,WN,1)
            x0=X[1,:].reshape(1,NX)
            xlog.append(x0)
            
            q=quaternion_from_euler(0,0,x0[0,0])
            new_msg.points[p].transforms[0].rotation.x=q[0]
            new_msg.points[p].transforms[0].rotation.y=q[1]
            new_msg.points[p].transforms[0].rotation.z=q[2]
            new_msg.points[p].transforms[0].rotation.w=q[3]
            
            
            
            #publish marker array
            marker=Marker()
            marker.header.frame_id='world'
            marker.type=Marker.ARROW
            #marker.action=Marker.ADD
            waypoint=new_msg.points[p].transforms[0].translation
            marker.pose.position.x=waypoint.x
            marker.pose.position.y=waypoint.y
            marker.pose.position.z=waypoint.z
            marker.pose.orientation.x=q[0]
            marker.pose.orientation.y=q[1]
            marker.pose.orientation.z=q[2]
            marker.pose.orientation.w=q[3]
            
            pose=marker.pose
            pose_array.poses.append(pose)
            
                # Set the scale, color, and other properties of the marker
            marker.scale.x = 1.0
            marker.scale.y = 0.1
            marker.scale.z = 0.1
            marker.color.a = 1.0
            marker.color.r = 1.0
            marker.color.g = 0.0
            marker.color.b = 0.0

            marker_array.markers.append(marker)
            
            #marker_pub.publish(marker_array)
            #pose_pub.publish(pose_array)
        #xlog_arr=np.vstack(xlog)
        
        
    def pc_callback(self,pc_msg):
        #downsample by o3d
        cloud = o3d.geometry.PointCloud()
        pc_np=pc2_msg_to_np(pc_msg)
        cloud.points=o3d.utility.Vector3dVector(pc_np.T)
        
        
        target_size=pc_size-10
        sampling_ratio=float(target_size)/pc_np.shape[1]
        
        downsampled_cloud=cloud.uniform_down_sample(int(1/sampling_ratio))
        pc_downsamped_np=np.asarray(downsampled_cloud.points).T        
        
        np.save('pc_test.npy',pc_downsamped_np)
        #print(pc_downsamped_np.shape)
        
        pc_downsamped_msg=np_to_pc2_msg(pc_downsamped_np)
        pc_downsamped_msg.header.frame_id="world"
        
        self.pc_pub.publish(pc_downsamped_msg)
        
        self.pc_W[0,:]=self.target_position[0]
        self.pc_W[1,:]=self.target_position[1]
        self.pc_W[2,:]=self.target_position[2]
        self.pc_W[:,:20]=pc_downsamped_np[:,:20]
        
        
    def run(self):
        rate=rospy.Rate(10)
        while not rospy.is_shutdown():
            """_
            summary_
            """
            
            
if __name__ == '__main__':
    try:
        path_publisher = ViewPlanning()
        path_publisher.run()
    except rospy.ROSInterruptException:
        pass 